package org.example.weboj.handler;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

public class SignalingHandler extends TextWebSocketHandler {

    // sessionId -> WebSocketSession
    private static final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();

    // roomName -> Set of sessions
    private static final Map<String, Set<WebSocketSession>> rooms = new ConcurrentHashMap<>();

    // sessionId -> roomName
    private static final Map<String, String> sessionRoomMap = new ConcurrentHashMap<>();

    private static final ObjectMapper objectMapper = new ObjectMapper();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        sessions.put(session.getId(), session);
        System.out.println("User connected: " + session.getId());
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        JsonNode msg = objectMapper.readTree(message.getPayload());
        String type = msg.get("type").asText();

        if ("join".equals(type)) {
            String room = msg.get("room").asText();
            joinRoom(session, room);
        } else if ("leave".equals(type)) {
            leaveRoom(session);
        }
        else {
            // signaling message: broadcast to others in the same room
            String room = sessionRoomMap.get(session.getId());
            if (room == null) return;

            for (WebSocketSession peer : rooms.get(room)) {
                if (!peer.getId().equals(session.getId()) && peer.isOpen()) {
                    peer.sendMessage(message);
                }
            }
        }
    }

    private void leaveRoom(WebSocketSession session) {
        String room = sessionRoomMap.remove(session.getId());
        if (room != null) {
            Set<WebSocketSession> roomSessions = rooms.get(room);
            if (roomSessions != null) {
                roomSessions.remove(session);
                if (roomSessions.isEmpty()) {
                    rooms.remove(room);
                }
            }
        }

        System.out.println("User left room: " + session.getId());
    }

    private void joinRoom(WebSocketSession session, String roomName) throws Exception {
        rooms.putIfAbsent(roomName, ConcurrentHashMap.newKeySet());
        Set<WebSocketSession> room = rooms.get(roomName);

        if (room.size() >= 2) {
            // 房间满了
            ObjectNode errorMsg = objectMapper.createObjectNode();
            errorMsg.put("type", "error");
            errorMsg.put("message", "Room is full");
            session.sendMessage(new TextMessage(errorMsg.toString()));
            return;
        }

        room.add(session);
        sessionRoomMap.put(session.getId(), roomName);

        // 通知加入成功
        ObjectNode successMsg = objectMapper.createObjectNode();
        successMsg.put("type", "joined");
        successMsg.put("room", roomName);
        successMsg.put("id", session.getId());
        session.sendMessage(new TextMessage(successMsg.toString()));
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        sessions.remove(session.getId());

        String room = sessionRoomMap.remove(session.getId());
        if (room != null) {
            Set<WebSocketSession> roomSessions = rooms.get(room);
            if (roomSessions != null) {
                roomSessions.remove(session);
                if (roomSessions.isEmpty()) {
                    rooms.remove(room);
                }
            }
        }

        System.out.println("User disconnected: " + session.getId());
    }
}
